/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: ffts plus task dispatcher
 * Author: lilianlin
 * Create: 2023-05-01
 */

#ifndef HCCL_DISPATCHER_FFTS_PUB_H
#define HCCL_DISPATCHER_FFTS_PUB_H

#include <vector>
#include "dispatcher_pub.h"
#include "ffts_ctx_provider.h"

namespace hccl {
constexpr u16 CONTEXT_MAX_NUM = 128;
// 910B和910_93 qos值默认是6
constexpr u32 QOS_DEFAULT_VALUE = 6;

constexpr uint32_t SDMA_FP32_ATOMIC_ADD_SQE = 0x1E71;
constexpr uint32_t SDMA_FP32_ATOMIC_MOVE_SQE = 0x1E70;
constexpr uint16_t FFTS_TIMEOUT_MAX = 65535;

using FftsSdmaSqeHeader = union {
    uint32_t word;
    struct {
        uint32_t opcode : 4;   // 0:非atomic搬运 1:atomic add 2:atomic max 3:atomic min 4:atomic equal 5:memory set
                               // 6:L2Cache Preload 7:L2Cache Prewriteback 8:L2Cache invalid 9:L2Cache Flush
        uint32_t datatype : 4; // 0:int8 1:int16 2:int32 6:fp16 normal 7:fp32 8:bf16 9:fp16 sat
        uint32_t ie2 : 1;      // 完成后是否上报中断; HCCL设为0
        uint32_t sssv : 1;     // source address对应的substream id是否有效; HCCL设为1
        uint32_t dssv : 1;     // destination address对应的substream id是否有效; HCCL设为1
        uint32_t sns : 1;      // source address对应的安全属性, 0:secure 1:non-secure; HCCL设为1
        uint32_t dns : 1;      // destination address对应的安全属性, 0:secure 1:non-secure; HCCL设为1
        uint32_t qos : 4;      // 0
        uint32_t sro : 1;      // 0
        uint32_t dro : 1;      // 0
        uint32_t partid : 8;   // 0
        uint32_t mpam : 1;     // 0
        uint32_t pmg : 2;      // 0
        uint32_t format : 1;   // 0
        uint32_t res6 : 1;     // 0
    } bit;
};


class DispatcherFFTS : public DispatcherPub {
public:
    explicit DispatcherFFTS(const s32 deviceLogicId);
    ~DispatcherFFTS() override;

    HcclResult SignalRecord(HcclRtNotify signal, hccl::Stream &stream, u32 userRank, u64 offset = INVALID_U64,
        s32 stage = INVALID_VALUE_STAGE, bool inchip = false, u64 signalAddr = INVALID_U64,
        u32 notifyId = INVALID_UINT) override;
    HcclResult SignalWait(HcclRtNotify signal, hccl::Stream &stream, u32 userRank, u32 remoteUserRank,
        s32 stage = INVALID_VALUE_STAGE, bool inchip = false, u32 notifyId = INVALID_UINT,
        u32 timeOut = NOTIFY_INVALID_WAIT_TIME) override;
    HcclResult MemcpyAsync(hccl::DeviceMem &dst, const hccl::DeviceMem &src, hccl::Stream &stream,
        u32 remoteUserRank = INVALID_VALUE_RANKID, hccl::LinkType inLinkType = hccl::LinkType::LINK_ONCHIP) override;
    HcclResult ReduceAsync(const void *src, void *dst, u64 dataCount, const HcclDataType datatype, HcclReduceOp redOp,
        Stream &stream, HcclReduceType reduceType = HcclReduceType::HCCL_TBE_REDUCE) override;
    HcclResult InlineReduceAsync(const void *src, u64 dataCount, const HcclDataType datatype, HcclReduceOp redOp,
        Stream &stream, void *dst, u32 remoteUserRank = INVALID_VALUE_RANKID,
        hccl::LinkType inLinkType = hccl::LinkType::LINK_ONCHIP) override;
    HcclResult RdmaSend(u32 dbindex, u64 dbinfo, const struct send_wr &wr, hccl::Stream &stream,
        u32 remoteUserRank = INVALID_VALUE_RANKID) override;
    HcclResult RdmaSend(u32 dbindex, u64 dbinfo, const struct send_wr &wr, hccl::Stream &stream, u32 userRank,
        u64 offset) override;

    HcclResult ResetFftsCtx(bool enableCache, const std::string &key) override;
    HcclResult LaunchTasksEx(Stream &stream, std::vector<Stream> &subStreams) override;
    void JudgeFftsCtxInitialized(bool &fftsCtxInitFlag) override;
    void SetNormalMode() override;

private:
    HcclResult TbeReduceAsync(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
        HcclReduceOp redOp, Stream &stream, const void *dst);
    HcclResult VectorReduce(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
        HcclReduceOp redOp, Stream &stream, const void *dst);
    HcclResult VectorReduceLoop(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
        HcclReduceOp redOp, Stream &stream, const void *dst);
    HcclResult TailVectorReduce(const void *tailSrc1, const void *tailSrc2, u64 tailCount, const HcclDataType dataType,
        HcclReduceOp redOp, Stream &stream, void *tailDst);
    HcclResult InitFftsDescVectorReduce(int count, void *addrListDevMemPtr, void *funcAddr, uint32_t blockDim,
        Stream &stream);
    HcclResult ConstructFftsAicAivCtx(int count, void *addrListDevMemPtr, void *funcAddr, uint32_t blockDim,
        rtFftsPlusAicAivCtx_t *ctx, bool isInit);
    HcclResult RefreshFftsDescVectorReduce(int count, void *addrListDevMemPtr, void *funcAddr, uint32_t blockDim,
        Stream &stream);
    void EnsureFftsContextsSize();
    HcclResult InitFftsDescNotifyRecordRemote(HcclRtNotify signal, Stream &stream, u64 &signalAddr);
    HcclResult InitFftsDescNotifyWait(HcclRtNotify signal, Stream &stream);
    HcclResult InitFftsDescSdma(void *dst, const void *src, u64 cnt, Stream &stream, u32 sdmaSqeHeader);
    HcclResult InitFftsDescMemcpy(DeviceMem &dst, const DeviceMem &src, Stream &stream);
    HcclResult InitFftsDescSdmaReduce(void *dst, const void *src, u64 cnt, const HcclDataType &datatype,
        const HcclReduceOp &redOp, Stream &stream);
    HcclResult RefreshFftsDescSdma(void *dst, const void *src, u64 cnt, Stream &stream, u32 sdmaSqeHeader);
    HcclResult RefreshFftsDescSdmaReduce(void *dst, const void *src, u64 cnt, const HcclDataType &datatype,
        const HcclReduceOp &redOp, Stream &stream);
    HcclResult RefreshFftsDescMemcpy(DeviceMem &dst, const DeviceMem &src, Stream &stream);
    HcclResult InitFftsDescRdmaSend(u32 dbindex, u64 dbinfo, Stream &stream);
    HcclResult RefreshFftsDescRdmaSend(u32 dbindex, u64 dbinfo, Stream &stream);
    HcclResult UpdataFftsDescRelationship(Stream &stream);
    HcclResult InitFftsSuccessorStart(HcclRtNotify signal, Stream &stream);
    HcclResult InitFftsSuccessorEnd(HcclRtNotify signal, Stream &stream);
    inline void SkipFftsRefresh()
    {
        fftsCtxsPtr->refreshIndex++;
        return;
    }
    inline bool FftsCtxReady()
    {
        return fftsCtxsPtr->completed;
    }
    HcclResult GetSdmaSqeHeader(const HcclDataType &datatype, const HcclReduceOp &redOp,
        FftsSdmaSqeHeader &sdmaSqeHeader);
    HcclResult PrintFFTSDebugDetails(rtFftsPlusSqe_t &fftsPlusSqe, rtFftsPlusTaskInfo_t &task);
    HcclResult ConstructFftsSqe(rtFftsPlusSqe_t &fftsPlusSqe);
    HcclResult ConstructFftsTask(rtFftsPlusTaskInfo_t &task, rtFftsPlusSqe_t &fftsPlusSqe);
    HcclResult RdmaSendInternal(u32 dbindex, u64 dbinfo, hccl::Stream &stream);
    HcclResult ConstructFftsWriteValueCtx(u32 dbindex, u64 dbinfo, rtFftsPlusWriteValueCtx_t *ctx);
    HcclResult ConstructFftsSdmaCtx(void *dst, const void *src, u64 cnt, u32 sdmaSqeHeader, rtFftsPlusSdmaCtx_t *ctx,
        bool isInit);
    HcclResult ConstructFFtsNotifyWaitCtx(HcclRtNotify &signal, rtFftsPlusNotifyCtx_t *ctx);
    HcclResult ConstructFftsNotifyRecordRemoteCtx(HcclRtNotify &signal, rtFftsPlusWriteValueCtx_t *ctx,
        u64 &signalAddr);
    HcclResult ConstructFftsNotifyCtx(u32 notifyID, uint16_t contextType, rtFftsPlusNotifyCtx_t *ctx);
    bool IsInvalidRdmaParam(u32 dbindex, u64 dbinfo);
    __inline__ __attribute__((always_inline)) HcclResult SignalTaskParaSave(HcclRtNotify signal, Stream &stream,
        u32 userRank, u32 remoteUserRank, u64 offset, s32 stage, TaskType taskType, uint64_t beginTime);
    HcclResult SetFftsTailVectorReduceDescSdma(void *devMem, const void *tailSrc, u64 count,
        const HcclDataType dataType, HcclReduceOp redOp, Stream &stream);
    HcclResult SetFftsDescVectorReduce(const void *src, const void *dst, int count, void *addrListDevMemPtr,
        void *funcAddr, uint32_t blockDim, const HcclDataType dataType, HcclReduceOp redOp, Stream &stream);
    HcclFftsContextsInfo *fftsCtxsPtr;
    FftsCtxProvider fftsCtxProvider;
    std::vector<void*> argsHandleList;
    bool disableFfts_;
};
} // namespace hccl
#endif // HCCL_DISPATCHER_FFTS_PUB_H
